use std::{
    io::{self, Read, Write},
    iter,
};

use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};

const PAYLOAD_LEN_U16: u8 = 126;
const PAYLOAD_LEN_U64: u8 = 127;

#[derive(Debug, PartialEq, Clone)]
#[allow(dead_code)]
pub enum OpCode {
    TextFrame = 1,
    BinaryFrame = 2,
    ConnectionClose = 8,
    Ping = 9,
    Pong = 0xA,
}

impl From<u8> for OpCode {
    fn from(value: u8) -> Self {
        match value {
            1 => OpCode::TextFrame,
            2 => OpCode::BinaryFrame,
            8 => OpCode::ConnectionClose,
            9 => OpCode::Ping,
            0xA => OpCode::Pong,
            _ => unreachable!(),
        }
    }
}
#[derive(Debug)]
#[allow(dead_code)]
pub struct FrameHeader {
    pub fin: bool,
    pub rsv1: bool,
    pub rsv2: bool,
    pub rsv3: bool,
    pub opcode: OpCode,
    pub masked: bool,
    pub payload_length: u8,
}

impl FrameHeader {
    fn new(len: usize, opcode: OpCode) -> Self {
        let payload_length = if len < (PAYLOAD_LEN_U16 as usize) {
            len as u8
        } else if len < (u16::MAX as usize) {
            PAYLOAD_LEN_U16
        } else {
            PAYLOAD_LEN_U64
        };

        Self {
            fin: true,
            rsv1: false,
            rsv2: false,
            rsv3: false,
            masked: false,
            payload_length,
            opcode,
        }
    }

    fn serialize(&self) -> u16 {
        let op = self.opcode.clone() as u8;
        let b1 = ((self.fin as u8) << 7)
            | ((self.rsv1 as u8) << 6)
            | ((self.rsv2 as u8) << 5)
            | ((self.rsv3 as u8) << 4)
            | (op & 0x0F);

        let b2 = ((self.masked as u8) << 7) | ((self.payload_length as u8) & 0x7F);

        ((b1 as u16) << 8) | (b2 as u16)
    }
}

impl From<u16> for FrameHeader {
    fn from(raw: u16) -> Self {
        let opcode = ((raw >> 8) as u8) & 0x0F;
        let opcode = OpCode::from(opcode);

        // 0010 0010 0000 0000

        // 0111 1111 => 0x7F
        // 0000 1111 => 0x0F
        // 1000 0000 => 0x80
        // 0100 0000 => 0x40
        // 0010 0000 => 0x20
        // 0001 0000 => 0x10
        Self {
            fin: (raw >> 8) & 0x80 == 0x80,
            rsv1: (raw >> 8) & 0x40 == 0x40,
            rsv2: (raw >> 8) & 0x20 == 0x20,
            rsv3: (raw >> 8) & 0x10 == 0x10,
            opcode,
            masked: raw & 0x80 == 0x80,
            payload_length: (raw as u8) & 0x7F,
        }
    }
}

#[derive(Debug)]
#[allow(dead_code)]
pub struct Frame {
    pub header: FrameHeader,
    mask: Option<[u8; 4]>,
    pub payload: Vec<u8>,
}

impl From<&str> for Frame {
    fn from(payload: &str) -> Self {
        Self {
            header: FrameHeader::new(payload.len(), OpCode::TextFrame),
            payload: Vec::from(payload),
            mask: None,
        }
    }
}

impl Frame {
    pub fn is_close(&self) -> bool {
        self.header.opcode == OpCode::ConnectionClose
    }
    pub fn pong(ping_frame: &Self) -> Self {
        let payload = ping_frame.payload.clone();
        Self {
            header: FrameHeader::new(payload.len(), OpCode::Pong),
            payload: payload,
            mask: None,
        }
    }

    pub fn close(recv_frame: &Self) -> Self {
        let body = if recv_frame.payload.len() > 0 {
            let status_code = &recv_frame.payload[0..2];
            let mut body = Vec::with_capacity(2);
            let _ = body.write(status_code);
            body
        } else {
            Vec::new()
        };
        Self {
            header: FrameHeader::new(body.len(), OpCode::ConnectionClose),
            payload: body,
            mask: None,
        }
    }
    /// Read a frame from the given reader.
    /// see https://datatracker.ietf.org/doc/html/rfc6455#section-5.2
    pub fn from_read<T: Read>(value: &mut T) -> io::Result<Self> {
        let buf = value.read_u16::<BigEndian>()?;
        let header = FrameHeader::from(buf);
        let len = Self::read_length(header.payload_length, value)?;

        let mask_key = if header.masked {
            let mut buf = [0; 4];
            value.read_exact(&mut buf)?;
            Some(buf)
        } else {
            None
        };

        let mut payload: Vec<u8> = Vec::with_capacity(len);
        payload.extend(iter::repeat(0).take(len));
        value.read(&mut payload)?;

        if let Some(mask) = mask_key {
            Self::apply_mask(mask, &mut payload);
        }

        Ok(Self {
            header,
            mask: mask_key,
            payload,
        })
    }
    pub fn serialize(&self) -> Vec<u8> {
        let mut buf = vec![];

        let len = self.header.payload_length;

        let h: u16 = self.header.serialize();
        let _ = buf.write_u16::<BigEndian>(h);

        match len {
            PAYLOAD_LEN_U16 => buf
                .write_u16::<BigEndian>(self.payload.len() as u16)
                .unwrap(),
            PAYLOAD_LEN_U64 => buf
                .write_u64::<BigEndian>(self.payload.len() as u64)
                .unwrap(),
            _ => (),
        }

        buf.write(self.payload.as_slice()).unwrap();
        buf
    }

    fn apply_mask(mask: [u8; 4], payload: &mut [u8]) {
        for (idx, c) in payload.iter_mut().enumerate() {
            *c = *c ^ mask[idx % 4];
        }
    }

    fn read_length<T: Read>(payload_len: u8, input: &mut T) -> io::Result<usize> {
        return match payload_len {
            PAYLOAD_LEN_U64 => input
                .read_u64::<BigEndian>()
                .map(|v| v as usize)
                .map_err(From::from),
            PAYLOAD_LEN_U16 => input
                .read_u16::<BigEndian>()
                .map(|v| v as usize)
                .map_err(From::from),
            _ => Ok(payload_len as usize), // payload_len < 127
        };
    }
}
